package com.anivia.common.log;


import com.anivia.common.log.util.HttpWrapperUtil;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * ResponseWrapper
 *
 * @author wangj
 */
@Slf4j
public class AniviaResponseWrapper extends HttpServletResponseWrapper {

    private AniviaServletOutputStream aniviaServletOutputStream;

    private boolean getBody = false;

    public static AniviaResponseWrapper build(HttpServletResponse response) {
        if (response instanceof AniviaResponseWrapper) {
            return (AniviaResponseWrapper) response;
        } else {
            return new AniviaResponseWrapper(response);
        }
    }

    private AniviaResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        getBody = HttpWrapperUtil.isGetBody((HttpServletResponse) super.getResponse());
        if (getBody) {
            if (aniviaServletOutputStream == null) {
                aniviaServletOutputStream = new AniviaServletOutputStream();
            }
            return aniviaServletOutputStream;
        } else {
            return super.getOutputStream();
        }
    }

    public byte[] getResponseBody() {
        if (getBody) {
            return aniviaServletOutputStream.getBody();
        } else {
            return new byte[0];
        }
    }

    public void responseOutWrite() throws IOException {
        if (getBody) {
            responseOutWrite(aniviaServletOutputStream.getBody());
        }
    }

    public void responseOutWrite(byte[] body) throws IOException {
        if (getBody) {
            getResponse().getOutputStream().write(body);
        }
    }

    public void destroy() {
        try {
            if (aniviaServletOutputStream != null) {
                aniviaServletOutputStream.flush();
                aniviaServletOutputStream.close();
            }
        } catch (Exception e) {
            log.error("流关闭异常", e);
        }
    }

    public static class AniviaServletOutputStream extends ServletOutputStream {

        private final ByteArrayOutputStream output = new ByteArrayOutputStream();

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener listener) {

        }

        @Override
        public void write(int b) throws IOException {
            output.write(b);
        }

        @Override
        public void write(byte[] b) throws IOException {
            output.write(b);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            output.write(b, off, len);
        }

        public byte[] getBody() {
            return output.toByteArray();
        }

        @Override
        public void close() throws IOException {
            output.close();
        }
    }

}
